//////////////////////////////////////////////////////////////////
// (c) Copyright 2010-  by Jeongnim Kim
//////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////
//   National Center for Supercomputing Applications &
//   Materials Computation Center
//   University of Illinois, Urbana-Champaign
//   Urbana, IL 61801
//   e-mail: jnkim@ncsa.uiuc.edu
//
// Supported by 
//   National Center for Supercomputing Applications, UIUC
//   Materials Computation Center, UIUC
//////////////////////////////////////////////////////////////////
// -*- C++ -*-
/**@file fft3d_comm.cpp
 */
#include <config.h>
#include <string>
#include <iostream>
#include <iomanip>
#include <cstdlib>
#if (__GNUC__ == 2)
#include <strstream.h>
#else
#include <sstream>
#endif
#include <strstream>
#include <fstream>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Utilities/Timer.h>
//#define DOSTUFF

inline void print_help(const std::string& msg)
{
  printf("%s -s fft_dim -y mpigridY -i iterations \n",msg.c_str());
}

using namespace std;

namespace mpi
{

  template<typename T>
  struct PencilDecomp
  {
    /** enumeration for the data ordering */
    enum {PXYZ=0, PXZY, PYZX};
    /** enumerate row and column */
    enum {IROW=0, ICOL=1};
    //const int Index[3][3];
    //enum {I_XYZ_X=0, I_XYZ_Y=1, I_XYZ_Z=2};
    //enum {I_XZY_X=0, I_XZY_Y=2, I_XZY_Z=1};
    //enum {I_YZX_X=2, I_YZX_Y=0, I_YZX_Z=1};
    /** size of the problem Length^3 */
    int Length;
    /** save the rank */
    int Rank;
    /** processor grid */
    int pGrid[2];
    /** grid coordination */
    int pCoord[2];
    /** lower bound of the block owned */
    int blockMin[2];
    /** upper bound of the block owned */
    int blockMax[2];
    /** size of the blocks */
    int blockSize[2];
    /** size of the messages per node */
    size_t msgSize[2];
    /** comm for row */
    communicator *rowComm;
    /** comm for column */
    communicator *colComm;
    /** buffer to perform rowComm->alltoall */
    vector<T> rowBuffer;
    /** buffer to perform colComm->alltoall */
    vector<T> colBuffer;
    vector<int> rowCounts, rowOffsets;
    vector<int> colCounts, colOffsets;

    PencilDecomp(int l, int nx, int ny, int ipx, int ipy)
      :Length(l),Rank(ipy+ipx*ny),rowComm(0),colComm(0)
      {
        pGrid[0]=nx;pGrid[1]=ny;
        pCoord[0]=ipx;pCoord[1]=ipy;
        partition();
      }

    /** constructor
     * @param l Length
     * @param ny number of partitions in the column
     * @param c 
     */
    PencilDecomp(int l, int ny, communicator& c)
      :Length(l),rowComm(0),colComm(0)
      {
        Rank=c.rank();
        pGrid[0]=c.size()/ny;pGrid[1]=ny;
        pCoord[0]=Rank/ny;
        pCoord[1]=Rank%ny;
        setCommunicator(c);
        partition();
      }

    ~PencilDecomp()
    {
      if(rowComm) delete rowComm;
      if(colComm) delete colComm;
    }

    inline void partition()
    {
      partition(Length,pGrid[0],pCoord[0],blockSize,blockMin,blockMax,IROW);
      partition(Length,pGrid[1],pCoord[1],blockSize,blockMin,blockMax,ICOL);

      msgSize[IROW]=blockSize[IROW]*blockSize[ICOL]*Length/rowComm->size();
      msgSize[ICOL]=blockSize[IROW]*blockSize[ICOL]*Length/colComm->size();

      rowBuffer.resize(msgSize[IROW]*rowComm->size(),Rank);
      colBuffer.resize(msgSize[ICOL]*colComm->size(),Rank);

      rowCounts.resize(rowComm->size(),msgSize[IROW]);
      rowOffsets.resize(rowComm->size());
      for(int i=0; i<rowComm->size(); ++i) rowOffsets[i]=i*msgSize[IROW];

      colCounts.resize(colComm->size(),msgSize[ICOL]);
      colOffsets.resize(colComm->size());
      for(int i=0; i<colComm->size(); ++i) colOffsets[i]=i*msgSize[ICOL];
    }

    void setCommunicator(communicator& c)
    {
      if(rowComm) delete rowComm;
      if(colComm) delete colComm;
      rowComm= new communicator(c.split(pCoord[0],pCoord[1]));
      colComm= new communicator(c.split(pCoord[1],pCoord[0]));
    }

    void print(int plan, std::ostream& os)
    {
      char f[512];
      sprintf(f,"%6d %3d %3d %3d %3d %6d %6d %6d %6d %6d(%d) %6d(%d) "
          ,Rank, pCoord[0], pCoord[1], rowComm->rank(), colComm->rank()
          ,blockMin[IROW],blockMax[IROW]
          ,blockMin[ICOL],blockMax[ICOL]
          ,blockMax[IROW]-blockMin[IROW] ,blockSize[IROW]
          ,blockMax[ICOL]-blockMin[ICOL] ,blockSize[ICOL]
          );
      os << f;
    }

    inline void fft1d(int dir)
    {
#if defined(DOSTUFF)
      if(dir)
        spectrum();
      else
        rk4();
#endif
    }

    inline void spectrum()
    {
#if defined(DOSTUFF)
      std::copy(rowBuffer.begin(),rowBuffer.end(),colBuffer.begin());
#endif
    }

    inline void rk4()
    {
#if defined(DOSTUFF)
      std::copy(colBuffer.begin(),colBuffer.end(),rowBuffer.begin());
#endif
    }

    void fft3d(int var)
    {
      fft1d(0);
      alltoall(*rowComm,&rowBuffer[0],&colBuffer[0],msgSize[IROW]);
      fft1d(1);
      alltoall(*colComm,&colBuffer[0],&rowBuffer[0],msgSize[ICOL]);
      fft1d(2);
    }

    void fft3dV(int var)
    {
      fft1d(0);
      alltoallv(*rowComm,&rowBuffer[0],&rowCounts[0],&rowOffsets[0],&colBuffer[0],&rowCounts[0],&rowOffsets[0]);
      fft1d(1);
      alltoallv(*colComm,&colBuffer[0],&colCounts[0],&colOffsets[0],&rowBuffer[0],&colCounts[0],&colOffsets[0]);
      fft1d(2);
    }

    void print(std::ostream& os)
    {
      os << "XYZ "; print(PXYZ,os); 
      os << " size= " << rowBuffer.size() 
        << " msgSize " << msgSize[IROW]*sizeof(T) << " " << msgSize[ICOL]*sizeof(T) << endl;
      //os << "XZY "; print(P_XZY,os);
      //os << "YZX "; print(P_YZX,os);
      /*
      const int nb=2;
      {
        std::vector<double> v(rowComm->size()*nb);
        std::vector<double> vout(rowComm->size()*nb,-1);
        for(int i=0,ii=0; i< rowComm->size(); ++i)
          for(int k=0;k<nb;++k, ++ii) v[ii]=Rank*nb+k;
        alltoall(*rowComm,&v[0],&vout[0],nb);

        os << "Row-In ";
        for(int i=0; i<v.size(); ++i) os << setw(4) << v[i];
        os << endl;

        os << "Row-Out ";
        for(int i=0; i<vout.size(); ++i) os << setw(4) << vout[i];
        os << endl;
      }

      {
        std::vector<double> v(colComm->size()*nb);
        std::vector<double> vout(colComm->size()*nb,-1);
        for(int i=0,ii=0; i< colComm->size(); ++i)
          for(int k=0;k<nb;++k, ++ii) v[ii]=Rank*nb+k;
        alltoall(*colComm,&v[0],&vout[0],nb);

        os << "Col-In ";
        for(int i=0; i<v.size(); ++i) os << setw(4) << v[i];
        os << endl;

        os << "Col-Out ";
        for(int i=0; i<vout.size(); ++i) os << setw(4) << vout[i];
        os << endl;
      }
      */
      //os << "Reduce " << v << " " << vr << " " << vc << std::endl;
    }

    /** partition a vector n by p
     * @param n dimension in 1D
     * @param p number of partitions
     * @param ip rank in p
     * @param bsize size of the block (n%p)? n/p+1 :n/p;
     * @param bmin lower bound
     * @param bmax upper bound
     * @param index for the plan
     */
    inline void partition(int n, int p, int ip, int *bsize, int *bmin, int *bmax, int index)
    {
      int np=n/p;//n per partition
      int p_more=n-p*np;//partitions n/p
      int p_less=p-p_more;//partitions with n/p+1
      bsize[index]=(p_more)?np+1:np;//size
      if(ip<p_less)
      {
        bmin[index]=np*ip;
        bmax[index]=bmin[index]+np;
      }
      else
      {
        bmin[index]=(np+1)*(ip-p_less)+np*p_less;
        bmax[index]=bmin[index]+np+1;
      }
    }
  };
}

int main(int argc, char** argv)
{

  mpi::environment env(argc,argv);
  mpi::communicator mycomm;

  int fftsize=786;
  int npy=16;
  int niters=5;

  bool debug=false;
  int opt;
  while((opt = getopt(argc, argv, "hds:y:i:")) != -1) {
    switch(opt) {
      case 'h':
        print_help("Help Message");
        return 1;
      case 'd':
        debug=true;
        break;
      case 's'://problem size
        fftsize=atoi(optarg);
        break;
      case 'y': //processor grid
        npy=atoi(optarg);
        break;
      case 'i': //processor grid
        niters=atoi(optarg);
        break;
      default:
        print_help("Unknown options");
        return 1;
    }
  }

  //create 2d decomp
  mpi::PencilDecomp<double> decomp(fftsize,npy,mycomm);
  //decomp.print(cout);
  //if(debug)
  //{
  //  char fname[64];
  //  sprintf(fname,"debug.p%04d",mycomm.rank());
  //  ofstream fout(fname);
  //  decomp.print(fout);
  //}

  decomp.fft3d(0);
  decomp.fft3d(1);

  qmcplusplus::Timer c;
  for(int iter=0; iter<niters; ++iter)
  {
    for(int iv=0; iv<6; ++iv) decomp.fft3d(iv);
    decomp.spectrum();
    for(int iv=0; iv<3; ++iv) decomp.fft3d(iv);
    decomp.rk4();
  }
  double dt=c.elapsed();
  mpi::reduce(mycomm,dt);

  if(mycomm.rank()==0) 
  {
    printf("#fftsize tasks npy msgYZ msg XY elapsed\n");
    printf("A2A %d %d %d %d %d %12.4e\n"
        ,fftsize,mycomm.size(),npy
        ,static_cast<int>(decomp.msgSize[0]*sizeof(double)), static_cast<int>(decomp.msgSize[1]*sizeof(double))
        ,dt/static_cast<double>(mycomm.size()*niters));
    }

  //int npx=mycomm.size()/npy;
  //int i=mycomm.rank()/npy;
  //int j=mycomm.rank()%npy;
  //PencilDecomp decomp(fftsize,npx,npy,i,j);
  //decomp.setCommunicator(mycomm);

  return 0;
}

/***************************************************************************
 * $RCSfile$   $Author: jnkim $
 * $Revision: 1770 $   $Date: 2007-02-17 17:45:38 -0600 (Sat, 17 Feb 2007) $
 * $Id: fft2d.cpp 1770 2007-02-17 23:45:38Z jnkim $ 
 ***************************************************************************/
